#!/usr/bin/python
#
# Copyright (C) 2012-2014 qar
# License: http://www.gnu.org/licenses/gpl.html GPL version 2
#

import os
import subprocess
import sys
import time

class CpuStats:
  def __init__(self):
    self.oldCpuStats = []
    self.hyperThreading = CpuStats.getHyperThreading()

  @staticmethod
  def getCpuStats():
    result = []
    with open("/proc/stat", "r") as f:
      for line in f:
        if line.startswith("cpu "):
          tokens = line.split()
          if len(tokens) == 11:
            for t in tokens[1:]:
              try:
                result.append(float(t))
              except:
                result.append(0.00001)
    return result

  @staticmethod
  def getHyperThreading():
    try:
      with open("/sys/devices/system/cpu/cpu0/topology/thread_siblings_list", "r") as f:
        for line in f:
          if len(line.split(",")) > 1:
            return True
    except:
      pass
    return False

  def getCpuLoad(self):
    newCpuStats = CpuStats.getCpuStats()
    load = 0.0
    if self.oldCpuStats:
      cpuDelta = [a - b for a, b in zip(newCpuStats, self.oldCpuStats)]
      load = 1.0 - cpuDelta[3] / (sum(cpuDelta) if sum(cpuDelta) != 0 else 1)
    self.oldCpuStats = newCpuStats
    if self.hyperThreading:
      load = min(1.0, load * 2.0)
    return load

def getFileSize(src):
  return int(subprocess.check_output(['du', '--block-size=1', '--summarize', src]).split('\t')[0])

def benchmarkArchiver(name, cmd, src, dst, disk):
  print >> sys.stderr, 'Benchmarking:', name
  subprocess.check_call(['rm', '-rf', dst])
  subprocess.check_call(['blockdev', '--flushbufs', disk])
  subprocess.check_call(['bash', '-c', 'echo 3 | tee /proc/sys/vm/drop_caches'])
  inputSizeB = getFileSize(src)
  inputSizeMB = inputSizeB / 1024.0 / 1024.0
  print >> sys.stderr, 'Input: {0}, size {1} MB ({2} B)'.format(src, inputSizeMB, inputSizeB)
  print >> sys.stderr, 'Executing:', cmd
  t1 = time.time()
  cpuStats = CpuStats()
  cpuStats.getCpuLoad()
  subprocess.check_call(cmd)
  t2 = time.time()
  cpuLoad = 100 * cpuStats.getCpuLoad()
  t = t2 - t1
  outputSizeB = getFileSize(dst)
  outputSizeMB = outputSizeB / 1024.0 / 1024.0
  ratio = outputSizeB * 100.0 / inputSizeB
  speed = inputSizeMB / t
  subprocess.check_call(['rm', '-rf', dst])
  print >> sys.stderr, 'Output: {0}, size {1} MB ({2} B)'.format(dst, outputSizeMB, outputSizeB)
  print >> sys.stderr, 'Ratio: {0} %'.format(ratio)
  print >> sys.stderr, 'Time: {0} s'.format(t)
  print >> sys.stderr, 'Speed: {0} MB/s'.format(inputSizeMB / t)
  print >> sys.stderr, 'CPU load: {0} %'.format(cpuLoad)
  print '|| {0} || {1} || {2:.0f} MB || {3:.0f} MB || {4}% || {5:.0f}s || {6:.0f} MB/s || {7:.0f}% ||'.format(name, ' '.join(cmd), inputSizeMB, outputSizeMB, int(ratio), t, speed, cpuLoad)

if not os.path.exists('linux-3.10'):
  if os.path.exists('linux-3.10.tar.xz'):
    subprocess.check_call(['rm', '-f', 'linux-3.10.tar.xz'])
  subprocess.check_call(['wget', 'https://www.kernel.org/pub/linux/kernel/v3.x/linux-3.10.tar.xz'])
  subprocess.check_call(['tar', '-xJf', 'linux-3.10.tar.xz'])
  subprocess.check_call(['rm', '-f', 'linux-3.10.tar.xz'])

benchmarkArchiver('qar store', ['qar', '-z', '0', 'linux-3.10'], 'linux-3.10', 'linux-3.10.qar', '/dev/sda')
benchmarkArchiver('qar', ['qar', 'linux-3.10'], 'linux-3.10', 'linux-3.10.qar', '/dev/sda')
benchmarkArchiver('tar store', ['tar', '-cf', 'linux-3.10.tar', 'linux-3.10'], 'linux-3.10', 'linux-3.10.tar', '/dev/sda')
benchmarkArchiver('tar gz', ['tar', '-czf', 'linux-3.10.tar.gz', 'linux-3.10'], 'linux-3.10', 'linux-3.10.tar.gz', '/dev/sda')
benchmarkArchiver('tar bz', ['tar', '-cjf', 'linux-3.10.tar.bz', 'linux-3.10'], 'linux-3.10', 'linux-3.10.tar.bz', '/dev/sda')
benchmarkArchiver('tar xz', ['tar', '-cJf', 'linux-3.10.tar.xz', 'linux-3.10'], 'linux-3.10', 'linux-3.10.tar.xz', '/dev/sda')
benchmarkArchiver('zip store', ['zip', '-0', '-q', '-r', 'linux-3.10.zip', 'linux-3.10'], 'linux-3.10', 'linux-3.10.zip', '/dev/sda')
benchmarkArchiver('zip', ['zip', '-q', '-r', 'linux-3.10.zip', 'linux-3.10'], 'linux-3.10', 'linux-3.10.zip', '/dev/sda')

if not os.path.exists('small-10k'):
  if not os.path.exists('pg844.txt'):
    subprocess.check_call(['wget', 'http://www.gutenberg.org/cache/epub/844/pg844.txt'])
  subprocess.check_call(['mkdir', 'small-10k'])
  subprocess.check_call(['bash', '-c', 'cd small-10k && for i in {1..10000} ; do cp ../pg844.txt ${i}.txt ; done'])

benchmarkArchiver('qar store', ['qar', '-z', '0', 'small-10k'], 'small-10k', 'small-10k.qar', '/dev/sda')
benchmarkArchiver('qar', ['qar', 'small-10k'], 'small-10k', 'small-10k.qar', '/dev/sda')
benchmarkArchiver('tar store', ['tar', '-cf', 'small-10k.tar', 'small-10k'], 'small-10k', 'small-10k.tar', '/dev/sda')
benchmarkArchiver('tar gz', ['tar', '-czf', 'small-10k.tar.gz', 'small-10k'], 'small-10k', 'small-10k.tar.gz', '/dev/sda')
benchmarkArchiver('tar bz', ['tar', '-cjf', 'small-10k.tar.bz', 'small-10k'], 'small-10k', 'small-10k.tar.bz', '/dev/sda')
benchmarkArchiver('tar xz', ['tar', '-cJf', 'small-10k.tar.xz', 'small-10k'], 'small-10k', 'small-10k.tar.xz', '/dev/sda')
benchmarkArchiver('zip store', ['zip', '-0', '-q', '-r', 'small-10k.zip', 'small-10k'], 'small-10k', 'small-10k.zip', '/dev/sda')
benchmarkArchiver('zip', ['zip', '-q', '-r', 'small-10k.zip', 'small-10k'], 'small-10k', 'small-10k.zip', '/dev/sda')
